#!/bin/bash
# LingShu消息队列 - Docker健康检查脚本

set -e

# 配置
LINGSHU_HOST="${LINGSHU_HOST:-localhost}"
LINGSHU_PORT="${LINGSHU_PORT:-4222}"
LINGSHU_HTTP_PORT="${LINGSHU_HTTP_PORT:-8222}"
TIMEOUT="${HEALTHCHECK_TIMEOUT:-5}"

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

# 日志函数
log_info() {
    echo -e "${GREEN}[HEALTH]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[HEALTH]${NC} $1"
}

log_error() {
    echo -e "${RED}[HEALTH]${NC} $1"
}

# 检查TCP端口连通性
check_tcp_port() {
    local host="$1"
    local port="$2"
    local service_name="$3"
    
    if timeout "$TIMEOUT" bash -c "echo >/dev/tcp/$host/$port" 2>/dev/null; then
        log_info "$service_name 端口 $port 连通正常"
        return 0
    else
        log_error "$service_name 端口 $port 连通失败"
        return 1
    fi
}

# 检查HTTP端点
check_http_endpoint() {
    local url="$1"
    local service_name="$2"
    
    if curl -f -s --max-time "$TIMEOUT" "$url" >/dev/null 2>&1; then
        log_info "$service_name HTTP端点正常: $url"
        return 0
    else
        log_error "$service_name HTTP端点异常: $url"
        return 1
    fi
}

# 检查进程状态
check_process() {
    local process_name="$1"
    
    if pgrep -f "$process_name" >/dev/null 2>&1; then
        log_info "进程 $process_name 运行正常"
        return 0
    else
        log_error "进程 $process_name 未运行"
        return 1
    fi
}

# 检查内存使用
check_memory_usage() {
    local max_memory_percent="${HEALTHCHECK_MAX_MEMORY:-90}"
    
    # 获取内存使用率
    local memory_usage
    memory_usage=$(free | grep Mem | awk '{printf "%.0f", $3/$2 * 100.0}')
    
    if [[ "$memory_usage" -lt "$max_memory_percent" ]]; then
        log_info "内存使用率正常: ${memory_usage}%"
        return 0
    else
        log_warn "内存使用率过高: ${memory_usage}% (阈值: ${max_memory_percent}%)"
        return 1
    fi
}

# 检查磁盘空间
check_disk_space() {
    local max_disk_percent="${HEALTHCHECK_MAX_DISK:-85}"
    local data_dir="${LINGSHU_DATA_DIR:-/var/lib/lingshu}"
    
    # 获取磁盘使用率
    local disk_usage
    disk_usage=$(df "$data_dir" | tail -1 | awk '{print $5}' | sed 's/%//')
    
    if [[ "$disk_usage" -lt "$max_disk_percent" ]]; then
        log_info "磁盘空间正常: ${disk_usage}%"
        return 0
    else
        log_warn "磁盘空间不足: ${disk_usage}% (阈值: ${max_disk_percent}%)"
        return 1
    fi
}

# 检查LingShu服务状态
check_lingshu_service() {
    local health_url="http://$LINGSHU_HOST:$LINGSHU_HTTP_PORT/health"
    
    # 尝试获取健康状态
    local response
    response=$(curl -f -s --max-time "$TIMEOUT" "$health_url" 2>/dev/null || echo "")
    
    if [[ -n "$response" ]]; then
        # 解析响应（简化版本）
        if echo "$response" | grep -q "healthy\|ok\|running"; then
            log_info "LingShu服务健康状态正常"
            return 0
        else
            log_warn "LingShu服务健康状态异常: $response"
            return 1
        fi
    else
        log_error "无法获取LingShu服务健康状态"
        return 1
    fi
}

# 检查JetStream状态
check_jetstream_status() {
    if [[ "${LINGSHU_JETSTREAM_ENABLED:-true}" == "true" ]]; then
        local jetstream_url="http://$LINGSHU_HOST:$LINGSHU_HTTP_PORT/jetstream"
        
        if curl -f -s --max-time "$TIMEOUT" "$jetstream_url" >/dev/null 2>&1; then
            log_info "JetStream状态正常"
            return 0
        else
            log_warn "JetStream状态异常"
            return 1
        fi
    else
        log_info "JetStream未启用，跳过检查"
        return 0
    fi
}

# 检查集群状态
check_cluster_status() {
    if [[ -n "${LINGSHU_CLUSTER_NAME}" ]]; then
        local cluster_url="http://$LINGSHU_HOST:$LINGSHU_HTTP_PORT/cluster"
        
        if curl -f -s --max-time "$TIMEOUT" "$cluster_url" >/dev/null 2>&1; then
            log_info "集群状态正常"
            return 0
        else
            log_warn "集群状态异常"
            return 1
        fi
    else
        log_info "集群模式未启用，跳过检查"
        return 0
    fi
}

# 执行完整健康检查
perform_health_check() {
    local exit_code=0
    local checks_passed=0
    local total_checks=0
    
    log_info "开始LingShu健康检查..."
    
    # 基础连通性检查
    ((total_checks++))
    if check_tcp_port "$LINGSHU_HOST" "$LINGSHU_PORT" "LingShu主服务"; then
        ((checks_passed++))
    else
        exit_code=1
    fi
    
    ((total_checks++))
    if check_tcp_port "$LINGSHU_HOST" "$LINGSHU_HTTP_PORT" "HTTP监控"; then
        ((checks_passed++))
    else
        exit_code=1
    fi
    
    # 进程检查
    ((total_checks++))
    if check_process "lingshu-server"; then
        ((checks_passed++))
    else
        exit_code=1
    fi
    
    # 资源使用检查
    ((total_checks++))
    if check_memory_usage; then
        ((checks_passed++))
    fi
    
    ((total_checks++))
    if check_disk_space; then
        ((checks_passed++))
    fi
    
    # 服务状态检查
    ((total_checks++))
    if check_lingshu_service; then
        ((checks_passed++))
    else
        exit_code=1
    fi
    
    # 可选功能检查
    ((total_checks++))
    if check_jetstream_status; then
        ((checks_passed++))
    fi
    
    ((total_checks++))
    if check_cluster_status; then
        ((checks_passed++))
    fi
    
    # 输出检查结果
    log_info "健康检查完成: $checks_passed/$total_checks 项检查通过"
    
    if [[ "$exit_code" -eq 0 ]]; then
        log_info "LingShu服务健康状态: 正常 ✓"
    else
        log_error "LingShu服务健康状态: 异常 ✗"
    fi
    
    return $exit_code
}

# 快速检查模式
quick_check() {
    log_info "执行快速健康检查..."
    
    # 只检查关键服务
    if check_tcp_port "$LINGSHU_HOST" "$LINGSHU_PORT" "LingShu主服务" && \
       check_process "lingshu-server"; then
        log_info "快速检查: 正常 ✓"
        return 0
    else
        log_error "快速检查: 异常 ✗"
        return 1
    fi
}

# 主函数
main() {
    case "${1:-full}" in
        "quick")
            quick_check
            ;;
        "full"|"")
            perform_health_check
            ;;
        "tcp")
            check_tcp_port "$LINGSHU_HOST" "$LINGSHU_PORT" "LingShu主服务"
            ;;
        "http")
            check_http_endpoint "http://$LINGSHU_HOST:$LINGSHU_HTTP_PORT/health" "LingShu HTTP"
            ;;
        "process")
            check_process "lingshu-server"
            ;;
        "memory")
            check_memory_usage
            ;;
        "disk")
            check_disk_space
            ;;
        "help")
            echo "LingShu健康检查脚本使用说明:"
            echo "  $0 [full]    - 执行完整健康检查（默认）"
            echo "  $0 quick     - 执行快速健康检查"
            echo "  $0 tcp       - 检查TCP端口连通性"
            echo "  $0 http      - 检查HTTP端点"
            echo "  $0 process   - 检查进程状态"
            echo "  $0 memory    - 检查内存使用"
            echo "  $0 disk      - 检查磁盘空间"
            echo "  $0 help      - 显示此帮助信息"
            exit 0
            ;;
        *)
            log_error "未知的检查类型: $1"
            echo "使用 '$0 help' 查看使用说明"
            exit 1
            ;;
    esac
}

# 执行主函数
main "$@"
